# -*- coding: utf-8 -*
import os
import pandas as pd
import numpy as np

from log import Logger
from catboost import CatBoostClassifier
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score
from typing import List

DATA_PATH = './ai/data'
DATA_LEN = 39                    # 共收集了39个系统性能数据
MODEL_SAVE_PATH = './model'
CATBOOSE_INFO_PATH = 'catboost_info'

train_file = ['idle.csv', 'mysql.csv', 'nginx.csv', 'postgresql.csv', 'redis.csv']
target_to_scene = {
    0: 'idle',
    1: 'mysql',
    2: 'nginx',
    3: 'postgresql',
    4: 'redis',
}


class AutoIdentifier:
    def __init__(self):
        self._X_test = None
        self._X_train = None
        self._y_test = None
        self._y_train = None
        self._model = None
        self._logger = Logger(__name__).get_logger()
        self._data = pd.DataFrame()
        self._labels = []
        self._state = False

    def _gen_data(self):
        """将csv数据集中的数据读取出来放入data中"""
        for index, file_name in enumerate(train_file):
            file_path = os.path.join(DATA_PATH, file_name)
            df = pd.read_csv(file_path)
            self._data = pd.concat([self._data, df], ignore_index=True)
            self._labels.extend([index] * len(df))

    def _split_data(self, test_size=0.1, random_state=42):
        """划分数据集和测试集"""
        self._X_train, self._X_test, self._y_train, self._y_test = train_test_split(self._data, self._labels, test_size=test_size,
                                                                                random_state=random_state)

    def _gen_model(self, verbose=10, random_seed=42, learning_rate=0.1, num_trees=100):
        """生成CatBoost分类器"""
        self._model = CatBoostClassifier(
            verbose=verbose,
            random_seed=random_seed,  # 随机种子
            learning_rate=learning_rate,  # 学习率
            num_trees=num_trees,  # 树的数量,
            loss_function="MultiClass"
        )

    def _train(self, is_save_model=True):
        """训练模型"""
        self._model.fit(self._X_train, self._y_train)
        if is_save_model:
            if not os.path.exists(MODEL_SAVE_PATH):
                os.makedirs(MODEL_SAVE_PATH)
                self._model.save_model(os.path.join(MODEL_SAVE_PATH, "auto_identify.cbm"))

    def test(self) -> float:
        """测试模型，返回准确率"""
        y_pred = self._model.predict(self._X_test)
        acc = accuracy_score(self._y_test, y_pred)
        return acc

    def detect(self, data: List[float]) -> str:
        """检测当前场景"""
        if len(data) != DATA_LEN:
            raise ValueError("The length of 'data' must be 39.")

        data = np.array(data).reshape(1, -1)

        if os.path.exists(MODEL_SAVE_PATH):
            if not self._state:
                self._logger.info("Use pretrained model")
                self._state = True
            pretrained = CatBoostClassifier()
            pretrained.load_model(os.path.join(MODEL_SAVE_PATH, 'auto_identify.cbm'))
            target = pretrained.predict(data)
        else:
            if not self._state:
                self._logger.info("No pretrained model. Start training model.")
                self._state = True
            self._gen_data()
            self._split_data()
            self._gen_model()
            self._train()
            target = self._model.predict(data)

        scene = target_to_scene[target[0][0]]
        return scene
